package com.xyh.ollama_test.service;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.xyh.ollama_test.enums.Department;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.JsonWithInt;
import io.qdrant.client.grpc.Points.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;


import java.util.Arrays;
import java.util.List;
import java.util.Map;

import java.util.stream.Collectors;
import java.util.stream.IntStream;

@Service
public class SearchInfo {
    private static final Logger logger = LoggerFactory.getLogger(SearchInfo.class);
    private final WebClient webClient;
    private final QdrantClient qdrantClient;
    private static final String COLLECTION_NAME = "text_embeddings";
    private static final String EMBEDDING_MODEL = "bge-m3";
    private static final String DEEPSEEK_MODEL = "deepseek-r1:1.5b"; // DeepSeek模型名称
    private final PermissionService permissionService;

    public SearchInfo(WebClient.Builder webClientBuilder, QdrantClient qdrantClient,PermissionService permissionService) {
        this.webClient = webClientBuilder.baseUrl("http://localhost:11434").build();
        this.qdrantClient = qdrantClient;
        this.permissionService = permissionService;
    }

    public Mono<String> searchAndGenerateResponse(String text,String userId) {
        if (text == null || text.trim().isEmpty()) {
            return Mono.just("输入文本不能为空");
        }

        logger.info("开始处理文本搜索请求: {}", text);

        return getTextEmbedding(text)
                .flatMap(embedding -> searchInQdrant(embedding,userId)
                        .flatMap(retrievedText -> generateResponseWithDeepSeek(text, retrievedText))
                )
                .onErrorResume(e -> {
                    logger.error("处理请求时出错", e);
                    return Mono.just("处理请求时出错: " + e.getMessage());
                });
    }

    private Mono<float[]> getTextEmbedding(String text) {
        return webClient.post()
                .uri("/api/embeddings")
                .bodyValue(Map.of(
                        "model", EMBEDDING_MODEL,
                        "prompt", text
                ))
                .retrieve()
                .bodyToMono(Map.class)
                .flatMap(response -> {
                    try {
                        @SuppressWarnings("unchecked")
                        List<Double> embeddingList = (List<Double>) response.get("embedding");
                        float[] floatArray = new float[embeddingList.size()];
                        for (int i = 0; i < embeddingList.size(); i++) {
                            floatArray[i] = embeddingList.get(i).floatValue();
                        }
                        System.out.println("floatArray ========= "+Arrays.toString(floatArray));
                        return Mono.just(floatArray);
                    } catch (Exception e) {
                        logger.error("解析嵌入向量时出错", e);
                        return Mono.error(new RuntimeException("解析嵌入向量时出错"));
                    }
                });
    }

    private Mono<String> searchInQdrant(float[] embedding, String userId) {
        try {
            // Get user department with null check
            Department userDept = permissionService.getUserDepartment(userId);
            if (userDept == null || userDept == Department.UNKNOWN) {
                return Mono.just("用户部门信息不存在或无效，无法查询");
            }

            List<Float> vector = floatArrayToFloatList(embedding);

            Filter.Builder filterBuilder = Filter.newBuilder();

            // For non-admin users, add department filter
            if (userDept != Department.ADMIN) {
                FieldCondition fieldCondition = FieldCondition.newBuilder()
                        .setKey("department")
                        .setMatch(Match.newBuilder()
                                .setKeyword(userDept.name())
                                .build())
                        .build();
                System.out.println("fieldCondition ======== "+fieldCondition);

                filterBuilder.addMust(Condition.newBuilder().setField(fieldCondition).build());
            }

            SearchPoints searchRequest = SearchPoints.newBuilder()
                    .setCollectionName(COLLECTION_NAME)
                    .addAllVector(vector)
                    .setLimit(3)
                    .setWithPayload(WithPayloadSelector.newBuilder().setEnable(true).build())
                    .setFilter(filterBuilder.build())
                    .build();

            List<ScoredPoint> results = qdrantClient.searchAsync(searchRequest).get();

            if (!results.isEmpty()) {
                ScoredPoint point = results.get(0);
                Map<String, JsonWithInt.Value> payload = point.getPayloadMap();
                if (payload.containsKey("text")) {
                    return Mono.just(payload.get("text").getStringValue());
                }
            }
            return Mono.just("未找到相关信息或您没有访问权限");
        } catch (Exception e) {
            logger.error("Qdrant搜索过程中出错", e);
            return Mono.error(new RuntimeException("搜索过程中出错: " + e.getMessage()));
        }
    }

    private Mono<String> generateResponseWithDeepSeek(String query, String context) {
        System.out.println("query ======= "+query);
        System.out.println("context ======== "+context);
        if (context == null || context.trim().isEmpty() || context =="未找到相关信息或您没有访问权限" || context=="用户部门信息不存在或无效，无法查询"){
            return Mono.just("未找到相关信息或您没有访问权限");
        }
        String prompt = String.format("基于以下上下文回答问题:\n%s\n\n问题: %s", context, query);

        return webClient.post()
                .uri("/api/generate")
                .bodyValue(Map.of(
                        "model", DEEPSEEK_MODEL,
                        "prompt", prompt,
                        "stream", false
                ))
                .retrieve()
                .bodyToMono(String.class)  // 先以字符串接收
                .flatMap(response -> {
                    try {
                        // 手动解析 JSON
                        ObjectMapper mapper = new ObjectMapper();
                        Map<String, Object> responseMap = mapper.readValue(response, Map.class);
                        return Mono.just(responseMap.get("response").toString());
                    } catch (Exception e) {
                        return Mono.error(e);
                    }
                })
                .onErrorResume(e -> {
                    logger.error("调用DeepSeek API出错", e);
                    return Mono.just("无法生成回答，请稍后再试");
                });
    }

    private List<Float> floatArrayToFloatList(float[] array) {
        return IntStream.range(0, array.length)
                .mapToObj(i -> array[i])
                .collect(Collectors.toList());
    }
}